"""PDF转DOCX转换器。"""

import json
import logging
import os
from multiprocessing import Pool, cpu_count
from time import perf_counter
from typing import AnyStr, IO, Union

import fitz
from docx import Document

from .page.Page import Page
from .page.Pages import Pages
from pdf2docx.common.printPro import printPro

# 检查PyMuPDF版本
# 1.19.0 <= v <= 1.23.8, or v>=1.23.16
v = list(map(int, fitz.VersionBind.split(".")))
if v < [1, 19, 0] or [1, 23, 8] < v < [1, 23, 16]:
    raise SystemExit(
        "1.19.0 <= PyMuPDF <= 1.23.8, or PyMuPDF>=1.23.16 is required for pdf2docx."
    )

# 日志配置
logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")


class Converter:
    """PDF到DOCX转换器。

    * 使用``PyMuPDF``读取PDF文件获取原始布局数据，逐页获取文本、图像、绘图及其属性，如边界框、字体、大小、图片宽度、高度等。
    * 在文档级别分析布局，如页眉、页脚和页边距。
    * 将页面布局解析为docx结构，如段落及其属性（缩进、间距、文本对齐）；表格及其属性（边框、底纹、合并）。
    * 最后，使用``python-docx``生成docx文件。
    """

    def __init__(
        self, pdf_file: str = None, password: str = None, stream: bytes = None
    ):
        """使用给定的pdf文件路径初始化fitz对象。

        参数:
            pdf_file (str): pdf文件路径。
            stream   (bytes): 内存中的pdf文件。
            password (str): 加密pdf的密码。如果未加密，默认为None。
        """
        # fitz对象
        self.filename_pdf = pdf_file
        self.password = str(password or "")

        if not pdf_file and not stream:
            raise ValueError("必须提供pdf_file或stream其中之一。")

        # 初始化fitz的doc对象
        if stream:
            self._fitz_doc = fitz.Document(stream=stream)

        else:
            self._fitz_doc = fitz.Document(pdf_file)

        # 初始化空页面容器
        self._pages = Pages()

    @property
    def fitz_doc(self):
        return self._fitz_doc

    @property
    def pages(self):
        return self._pages

    def close(self):
        self._fitz_doc.close()

    @property
    def default_settings(self):
        """默认解析参数。"""
        return {
            "debug": False,  # 如果为True则绘制布局
            "ocr": 0,  # ocr状态：0 - 不进行ocr；1 - 执行ocr；2 - 已ocr处理的pdf
            "ignore_page_error": True,  # 如果为True，则某页失败不会中断转换过程
            "multi_processing": False,  # 如果为True，则使用多进程转换页面
            "cpu_count": 0,  # 使用多进程转换页面时的CPU数量
            "min_section_height": 20.0,  # 有效部分的最小高度。
            "connected_border_tolerance": 0.5,  # 如果间隙小于此值，则两个边框相交
            "max_border_width": 6.0,  # 最大边框宽度
            "min_border_clearance": 2.0,  # 两个边框之间的最小允许间隙
            "float_image_ignorable_gap": 5.0,  # 如果相交超过此值，则为浮动图像
            "page_margin_factor_top": 0.5,  # [0,1] 按因子减少顶部边距
            "page_margin_factor_bottom": 0.5,  # [0,1] 按因子减少底部边距
            "shape_min_dimension": 2.0,  # 如果宽度和高度都低于此值，则忽略形状
            "max_line_spacing_ratio": 1.5,  # 最大行间距比率：行间距/行高
            "line_overlap_threshold": 0.9,  # [0,1] 如果与其他线的交叉超过此值，则删除线
            "line_break_width_ratio": 0.5,  # 如果行宽与整个布局边界框的比率低于此值，则换行
            "line_break_free_space_ratio": 0.1,  # 如果整行中的空闲空间比率超过此值，则换行
            "line_separate_threshold": 5.0,  # 如果x距离超过此值，则为两条单独的线
            "new_paragraph_free_space_ratio": 0.85,  # 如果空闲空间与行高的比率超过此值，则为新段落
            "lines_left_aligned_threshold": 1.0,  # 如果两行的d_x0低于此值(Pt)，则左对齐
            "lines_right_aligned_threshold": 1.0,  # 如果两行的d_x1低于此值(Pt)，则右对齐
            "lines_center_aligned_threshold": 2.0,  # 如果两行的中心差低于此值，则居中对齐
            "clip_image_res_ratio": 4.0,  # 剪裁页面图像时的分辨率比率（与72dpi相比）
            "min_svg_gap_dx": 15.0,  # 如果水平间隙小于此值，则合并相邻的矢量图形
            "min_svg_gap_dy": 2.0,  # 如果垂直间隙小于此值，则合并相邻的矢量图形
            "min_svg_w": 2.0,  # 如果边界框宽度小于此值，则忽略矢量图形
            "min_svg_h": 2.0,  # 如果边界框高度小于此值，则忽略矢量图形
            "extract_stream_table": False,  # 提取表格时不考虑流表
            "parse_lattice_table": True,  # 是否解析格子表，设置为False可能会破坏布局
            "parse_stream_table": True,  # 是否解析流表，设置为False可能会破坏布局
            "delete_end_line_hyphen": False,  # 删除行尾连字符
            "raw_exceptions": False,  # 不吞异常
            "list_not_table": True,  # 避免将项目符号列表视为表格。
        }

    # -----------------------------------------------------------------------
    # 解析过程：加载 -> 分析文档 -> 解析页面 -> 创建docx
    # -----------------------------------------------------------------------

    def parse(self, start: int = 0, end: int = None, pages: list = None, **kwargs):
        """分三步解析页面：
        * 使用``PyMuPDF``打开PDF文件
        * 分析整个文档，如页面部分、页眉/页脚和边距
        * 解析指定页面，如段落、图像和表格

        参数:
            start (int, optional): 要处理的第一页。默认为0，即第一页。
            end (int, optional): 要处理的最后一页。默认为None，即最后一页。
            pages (list, optional): 要解析的页面索引范围。默认为None。
            kwargs (dict, optional): 配置参数。
        """
        return (
            self.load_pages(start, end, pages)
            .parse_document(**kwargs)
            .parse_pages(**kwargs)
        )

    def load_pages(self, start: int = 0, end: int = None, pages: list = None):
        """转换过程的步骤1：使用``PyMuPDF``打开PDF文件，尤其是处理密码加密的文件。

        参数:
            start (int, optional): 要处理的第一页。默认为0，即第一页。
            end (int, optional): 要处理的最后一页。默认为None，即最后一页。
            pages (list, optional): 要解析的页面索引范围。默认为None。
        """
        logging.info(self._color_output("[1/4] 打开文档..."))

        # 加密的pdf？
        if self._fitz_doc.needs_pass:
            if not self.password:
                raise ConversionException(f"需要密码来打开 {self.filename_pdf}。")

            elif not self._fitz_doc.authenticate(self.password):
                raise ConversionException("密码不正确。")

        # 初始化空页面
        num = len(self._fitz_doc)
        
        self._pages.reset([Page(id=i, skip_parsing=True) for i in range(num)])

        # 设置要解析的页面
        page_indexes = self._page_indexes(start, end, pages, num)
        for i in page_indexes:
            self._pages[i].skip_parsing = False

        return self

    def parse_document(self, **kwargs):
        """转换过程的步骤2：分析整个文档，如页面部分、页眉/页脚和边距。"""
        logging.info(self._color_output("[2/4] 分析文档..."))

        self._pages.parse(self.fitz_doc, **kwargs)
        return self

    def parse_pages(self, **kwargs):
        """转换过程的步骤3：解析页面，如段落、图像和表格。"""
        logging.info(self._color_output("[3/4] 解析页面..."))

        pages = [page for page in self._pages if not page.skip_parsing]
        num_pages = len(pages)
        for i, page in enumerate(pages, start=1):
            pid = page.id + 1
            logging.info("(%d/%d) 页面 %d", i, num_pages, pid)
            try:
                page.parse(**kwargs)
            except Exception as e:
                if kwargs["raw_exceptions"]:
                    raise
                if not kwargs["debug"] and kwargs["ignore_page_error"]:
                    logging.error("由于解析页面错误，忽略页面 %d：%s", pid, e)
                else:
                    raise ConversionException(f"解析页面 {pid} 时出错：{e}")

        return self

    def make_docx(self, filename_or_stream=None, **kwargs):
        """转换过程的步骤4：使用转换后的页面创建docx文件。

        参数:
            filename_or_stream (str, file-like): 要写入的docx文件。
            kwargs (dict, optional): 配置参数。
        """
        logging.info(self._color_output("[4/4] 创建页面..."))

        # 检查已解析的页面
        parsed_pages = list(filter(lambda page: page.finalized, self._pages))
        if not parsed_pages:
            raise ConversionException("没有已解析的页面。请先解析页面。")

        if not filename_or_stream:
            if self.filename_pdf:
                filename_or_stream = f'{self.filename_pdf[0:-len(".pdf")]}.docx'
                # 删除已存在的文件
                if os.path.exists(filename_or_stream):
                    os.remove(filename_or_stream)
            else:
                raise ConversionException("请指定docx文件名或类文件对象来写入。")

        # 逐页创建
        docx_file = Document()
        num_pages = len(parsed_pages)
        for i, page in enumerate(parsed_pages, start=1):
            if not page.finalized:
                continue  # 忽略未解析的页面
            pid = page.id + 1
            logging.info("(%d/%d) 页面 %d", i, num_pages, pid)
            try:
                page.make_docx(docx_file)
            except Exception as e:
                if kwargs["raw_exceptions"]:
                    raise
                if not kwargs["debug"] and kwargs["ignore_page_error"]:
                    logging.error("由于创建页面错误，忽略页面 %d：%s", pid, e)
                else:
                    raise MakedocxException(f"创建页面 {pid} 时出错：{e}")

        # 保存docx
        docx_file.save(filename_or_stream)

    # -----------------------------------------------------------------------
    # 存储 / 恢复解析结果
    # -----------------------------------------------------------------------

    def store(self):
        """以字典格式存储已解析的页面。"""
        return {
            "filename": os.path.basename(self.filename_pdf),
            "page_cnt": len(self._pages),  # 所有页面的计数
            "pages": [
                page.store() for page in self._pages if page.finalized
            ],  # 仅已解析的页面
        }

    def restore(self, data: dict):
        """从解析结果中恢复页面。"""
        # 如有必要，初始化空页面
        if not self._pages:
            num = data.get("page_cnt", 100)
            self._pages.reset([Page(id=i, skip_parsing=True) for i in range(num)])

        # 恢复页面
        for raw_page in data.get("pages", []):
            idx = raw_page.get("id", -1)
            self._pages[idx].restore(raw_page)

    def serialize(self, filename: str):
        """将已解析的页面写入指定的JSON文件。"""
        with open(filename, "w", encoding="utf-8") as f:
            f.write(json.dumps(self.store(), indent=4))

    def deserialize(self, filename: str):
        """从指定的JSON文件加载已解析的页面。"""
        with open(filename, "r") as f:
            data = json.load(f)
        self.restore(data)

    # -----------------------------------------------------------------------
    # 高级方法，如转换、提取表格
    # -----------------------------------------------------------------------

    def debug_page(
        self,
        i: int,
        docx_filename: str = None,
        debug_pdf: str = None,
        layout_file: str = None,
        **kwargs,
    ):
        """解析、创建和绘制单个页面用于调试。

        参数:
            i (int): 要转换的页面索引。
            docx_filename (str): 要写入的docx文件名。
            debug_pdf (str): 存储布局信息的新pdf文件。默认添加前缀``debug_``。
            layout_file (str): 存储已解析布局数据的新json文件。默认为``layout.json``。
        """
        # 包括调试信息
        # 调试模式下的fitz对象：绘制页面布局
        # 此调试pdf的文件路径：demo.pdf -> debug_demo.pdf
        path, filename = os.path.split(self.filename_pdf)
        if not debug_pdf:
            debug_pdf = os.path.join(path, f"debug_{filename}")
        if not layout_file:
            layout_file = os.path.join(path, "layout.json")
        kwargs.update(
            {"debug": True, "debug_doc": fitz.Document(), "debug_filename": debug_pdf}
        )

        # 解析并创建docx
        self.convert(docx_filename, pages=[i], **kwargs)

        # 用于调试的布局信息
        self.serialize(layout_file)

    def convert(
        self,
        docx_filename: Union[str, IO[AnyStr]] = None,
        start: int = 0,
        end: int = None,
        pages: list = None,
        **kwargs,
    ):
        """将指定的PDF页面转换为docx文件。

        参数:
            docx_filename (str, file-like, optional): 要写入的docx文件。默认为None。
            start (int, optional): 要处理的第一页。默认为0，即第一页。
            end (int, optional): 要处理的最后一页。默认为None，即最后一页。
            pages (list, optional): 页面索引范围。默认为None。
            kwargs (dict, optional): 配置参数。默认为None。

        有关配置参数的详细信息，请参阅 :py:meth:`~pdf2docx.converter.Converter.default_settings`。

        .. 注意::
            如果``docx_file``为None，则将扩展名从``pdf``更改为``docx``。

        .. 注意::
            * 如果``--zero_based_index=True``（默认），则``start``和``end``从零开始计数。
            * 如果省略``start``，则从第一页开始。
            * 如果省略``end``，则以最后一页结束。

        .. 注意::
            ``pages``比``start``和``end``具有更高的优先级。``start``和``end``仅在省略``pages``时有效。

        .. 注意::
            多进程处理仅适用于由``start``和``end``指定的连续页面。
        """
        t0 = perf_counter()
        logging.info("开始转换 %s", self.filename_pdf)
        settings = self.default_settings
        settings.update(kwargs)

        # 输入检查
        if pages and settings["multi_processing"]:
            raise ConversionException(
                '多进程处理仅适用于由"start"和"end"指定的连续页面。'
            )

        # 逐页转换
        if settings["multi_processing"]:
            self._convert_with_multi_processing(docx_filename, start, end, **settings)
        else:
            self.parse(start, end, pages, **settings).make_docx(
                docx_filename, **settings
            )

        logging.info("在%.2f秒内完成。", perf_counter() - t0)

    def extract_tables(
        self, start: int = 0, end: int = None, pages: list = None, **kwargs
    ):
        """从指定的PDF页面提取表格内容。

        参数:
            start (int, optional): 要处理的第一页。默认为0，即第一页。
            end (int, optional): 要处理的最后一页。默认为None，即最后一页。
            pages (list, optional): 页面索引范围。默认为None。
            kwargs (dict, optional): 配置参数。默认为None。

        返回:
            list: 已解析表格内容的列表。
        """
        # 首先解析页面
        settings = self.default_settings
        settings.update(kwargs)
        self.parse(start, end, pages, **settings)

        # 获取已解析的表格
        tables = []
        for page in self._pages:
            if page.finalized:
                tables.extend(page.extract_tables(**settings))

        return tables

    def _convert_with_multi_processing(
        self, docx_filename: str, start: int, end: int, **kwargs
    ):
        """基于页面索引使用多进程解析和创建页面。

        参考:

            https://pymupdf.readthedocs.io/en/latest/faq.html#multiprocessing
        """
        # 为进程制作参数向量
        cpu = (
            min(kwargs["cpu_count"], cpu_count())
            if kwargs["cpu_count"]
            else cpu_count()
        )
        prefix = "pages"  # 每个进程写入已解析页面的json文件
        vectors = [
            (
                i,
                cpu,
                start,
                end,
                self.filename_pdf,
                self.password,
                kwargs,
                f"{prefix}-{i}.json",
            )
            for i in range(cpu)
        ]

        # 启动解析进程
        pool = Pool()
        pool.map(self._parse_pages_per_cpu, vectors, 1)

        # 恢复已解析的页面数据
        for i in range(cpu):
            filename = f"{prefix}-{i}.json"
            if not os.path.exists(filename):
                continue
            self.deserialize(filename)
            os.remove(filename)

        # 创建docx文件
        self.make_docx(docx_filename, **kwargs)

    @staticmethod
    def _parse_pages_per_cpu(vector):
        """渲染文档的页面范围。

        参数:
            vector (list): 包含所需参数的列表。
                * 0  : 当前进程的段号
                * 1  : CPU数量
                * 2,3: 要处理的整个页面范围
                * 4  : pdf文件名
                * 5  : 加密pdf的密码
                * 6  : 配置参数
                * 7  : 存储已解析结果的json文件名
        """
        # 重新创建参数
        idx, cpu, s, e, pdf_filename, password, kwargs, json_filename = vector

        # 打开pdf获取页面计数：所有页面暂时都标记为解析
        # 因为此刻不知道要解析哪些页面
        cv = Converter(pdf_filename, password)
        cv.load_pages()

        # 指定要处理的页面
        e = e or len(cv.fitz_doc)
        all_indexes = range(s, e)
        num_pages = len(all_indexes)

        # 此cpu处理的页面段
        m = int(num_pages / cpu)
        n = num_pages % cpu
        seg_size = m + int(idx < n)
        seg_from = (m + 1) * idx + min(n - idx, 0)
        seg_to = min(seg_from + seg_size, num_pages)
        page_indexes = [all_indexes[i] for i in range(seg_from, seg_to)]

        # 现在，标记正确的页面
        for page in cv.pages:
            page.skip_parsing = True
        for i in page_indexes:
            cv.pages[i].skip_parsing = False

        # 解析页面并序列化数据以供进一步处理
        cv.parse_document(**kwargs).parse_pages(**kwargs).serialize(json_filename)
        cv.close()

    @staticmethod
    def _page_indexes(start, end, pages, pdf_len):
        """解析参数。"""
        if pages:
            indexes = [int(x) for x in pages]
        else:
            end = end or pdf_len
            s = slice(int(start), int(end))
            indexes = range(pdf_len)[s]

        return indexes

    @staticmethod
    def _color_output(msg):
        return f"\033[1;36m{msg}\033[0m"


class ConversionException(Exception):
    pass


class MakedocxException(ConversionException):
    pass
